import os
import sys
import shutil
import json
import numpy as np
import glob
import cv2
from sklearn.model_selection import train_test_split
from os import getcwd

def Change2json(txt_dir):
    txt_dir = txt_dir + '/'
    # 获取所有txt文件
    files = sorted(glob.glob(os.path.join(txt_dir, '*.txt')))
    # 提取文件名（不包括扩展名）
    files = [os.path.splitext(os.path.basename(i))[0] for i in files]

    def convert(size, box):
        x = float(box[1])
        w = float(box[3])
        y = float(box[2])
        h = float(box[4])
        x = x * size[0]
        w = w * size[0]
        y = y * size[1]
        h = h * size[1]
        xmin = (x+1) * 2.0 - w
        xmax = (x+1) * 2.0 + w
        ymin = (y+1) * 2.0 - h
        ymax = (y+1) * 2.0 + h
        xmin = xmin / 2.0
        xmax = xmax / 2.0
        ymin = ymin / 2.0
        ymax = ymax / 2.0
        return (xmin, ymin, xmax, ymax)

    # 设置工作目录
    wd = "D:\\ToolsetCatalog\\txt_json"
    print(wd)

    # 支持的图片格式
    supported_img_formats = ['.bmp', '.jpg', '.png']

    for txt_file in files:
        # 确定实际存在的图片格式
        image_path = None
        for suffix in supported_img_formats:
            potential_image_path = os.path.abspath(txt_dir + txt_file + suffix)
            if os.path.exists(potential_image_path):
                image_path = potential_image_path
                break
        
        if not image_path:
            print(f"警告: {txt_file} 没有找到对应的图片文件.")
            continue

        # 准备文件路径
        txt_filename = os.path.abspath(txt_dir + txt_file + ".txt")
        json_filename = os.path.abspath(txt_dir + txt_file + ".json")
        imagefilename = os.path.basename(image_path)

        # 读取TXT文件
        with open(txt_filename, 'r') as f:
            txt_data = f.readlines()

        height, width = cv2.imread(image_path).shape[:2]

        # 构建JSON格式数据
        shapes = []
        for line in txt_data:
            data = line.split()
            bb = convert((width, height), data)
            shape = {
                "label": "qx",
                "points": [
                    [float(bb[0]), float(bb[1])],
                    [float(bb[2]), float(bb[3])]
                ],
                "group_id": None,
                "shape_type": "rectangle",
                "flags": {}
            }
            shapes.append(shape)

        data = {
            "version": "5.0.1",
            "flags": {},
            "shapes": shapes,
            "imagePath": imagefilename,
            "imageData": None,
            "imageHeight": height,
            "imageWidth": width
        }

        # 将JSON数据写入文件
        with open(json_filename, 'w') as f:
            json.dump(data, f, indent=2)

if __name__ == '__main__':
    # 文件夹目录
    txt_dir = sys.argv[1]
    Change2json(txt_dir)
    print('OK!')